# SPDX-License-Identifier: Apache-2.0
# (C) Copyright IBM Corp. 2024.
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

import os
import sys

from data_processing.utils import get_dpk_logger

from . import PipelinesUtils


logger = get_dpk_logger()


def run_test(pipeline_package_path: str, endpoint: str = "http://localhost:8080/"):
    """
    Upload and run a single pipeline

    :param pipeline_package_path: Local path to the pipeline package.
    :param endpoint: endpoint to kfp service.
    :return the pipeline name as it appears in the kfp GUI.
    """
    tmout: int = 800
    wait: int = 60
    file_name = os.path.basename(pipeline_package_path)
    pipeline_name = os.path.splitext(file_name)[0]
    utils = PipelinesUtils(host=endpoint)
    pipeline = utils.upload_pipeline(
        pipeline_package_path=pipeline_package_path,
        pipeline_name=pipeline_name,
    )
    if pipeline is None:
        return None
    experiment = utils.get_experiment_by_name()
    run_id = utils.start_pipeline(pipeline, experiment, params={})
    if run_id is None:
        return None
    error = utils.wait_pipeline_completion(run_id=run_id, timeout=tmout, wait=wait)
    if error is not None:
        # Execution failed
        logger.warning(f"Pipeline {pipeline_name} failed with error {error}")
        return None
    logger.info(f"Pipeline {pipeline_name} successfully completed")
    return pipeline_name

def _set_run_id(pipeline_package_path: str):
    """
    Assign a dummy run ID value for testing purposes. By default, this value
    is empty and is set by the user during runtime.

    :param pipeline_package_path: Local path to the pipeline package.
    """
    import yaml
    import uuid

    try:
        stream = open(pipeline_package_path, "r")
        docs = list(yaml.load_all(stream, yaml.FullLoader))
        for doc in docs:
            if "root" in doc:
                doc["root"]["inputDefinitions"]["parameters"]["ray_run_id_KFPv2"]["defaultValue"] = uuid.uuid4().hex
        with open(pipeline_package_path, "w") as outfile:
            yaml.dump_all(docs, outfile)
    except Exception as e:
        logger.error(f"Failed to update run id value, exception {e}")
        sys.exit(1)

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Run sanity test")
    parser.add_argument("-c", "--command", type=str, choices=["upload", "sanity-test"])
    parser.add_argument("-e", "--endpoint", type=str, default="http://localhost:8080/")
    parser.add_argument("-p", "--pipeline_package_path", type=str, default="")

    args = parser.parse_args()
    match args.command:
        case "upload":
            file_name = os.path.basename(args.pipeline_package_path)
            pipeline_name = os.path.splitext(file_name)[0]
            utils = PipelinesUtils(host=args.endpoint)
            pipeline = utils.upload_pipeline(
                pipeline_package_path=args.pipeline_package_path,
                pipeline_name=pipeline_name,
            )
            if pipeline is None:
                sys.exit(1)
        case "sanity-test":
            _set_run_id(args.pipeline_package_path)
            run = run_test(
                endpoint=args.endpoint,
                pipeline_package_path=args.pipeline_package_path,
            )
            if run is None:
                sys.exit(1)
        case _:
            logger.warning("Unsupported command")
            exit(1)
